package middlewarex

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	jsoniter "github.com/json-iterator/go"

	"zhonghui/console/warehouse/helperx"
	"zhonghui/console/warehouse/logx"
)

// 返回结构体
type responseBodyWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

// 重写返回函数
func (w responseBodyWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

// ExchangeLogger 记录所有请求和返回日志中间件
func ExchangeLogger(needReq bool, needResp bool, logHandler func([]byte, string, string, string, string, int64)) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 忽略一些URI的日志
		switch c.Request.URL.String() {
		case "/health":
			c.Next()
			return
		}

		// 记录请求时间
		c.Set("request_time", time.Now())

		clientIp := c.ClientIP()
		requestUrl := c.Request.URL.String()
		userAgent := c.Request.UserAgent()
		requestId := c.GetString("RequestId")
		headers := make([]string, 0)
		for k, v := range c.Request.Header {
			headers = append(headers, fmt.Sprintf("%s:%s", k, v))
		}
		headersStr := strings.Join(headers, " ")

		// 开启请求日志
		if needReq {
			// 读取buffer
			buf, _ := ioutil.ReadAll(c.Request.Body)
			// 复制buffer
			body := ioutil.NopCloser(bytes.NewBuffer(buf))
			// 重新将body放回context
			c.Request.Body = body
			// 日志记录size限制
			if len(buf) > helperx.KB*10 {
				buf = []byte("too large")
			}

			// 记录请求日志
			logx.Zap().Infow("api_request",
				"request_id", requestId,
				"url", requestUrl,
				"headers", headersStr,
				"body", string(buf),
				"user_agent", userAgent,
				"client_ip", clientIp,
			)
		}

		// 读取返回
		w := &responseBodyWriter{body: &bytes.Buffer{}, ResponseWriter: c.Writer}
		c.Writer = w

		// 下一步
		c.Next()

		// 开启返回日志
		if needResp {
			resp := w.body.Bytes()

			// 返回时间
			var timeDiff int64
			requestTime, ex := c.Get("request_time")
			if ex {
				if v, ok := requestTime.(time.Time); ok {
					timeDiff = time.Now().Sub(v).Milliseconds()
				}
			}

			// 如不传，则使用默认返回日志处理器
			if logHandler == nil {
				defaultLogHandler(resp, requestId, requestUrl, userAgent, clientIp, timeDiff)
			} else {
				logHandler(resp, requestId, requestUrl, userAgent, clientIp, timeDiff)
			}
		}
	}
}

// 默认的返回日志记录
// 只能检测 {"meta":{"code":"SUCCESS"}, "data": null} 这种格式的返回体
// 如是其他格式的可以自己实现本方法
func defaultLogHandler(resp []byte, requestId, requestUrl, userAgent, clientIp string, timeDiff int64) {
	// 解析返回中的code
	r := struct {
		Meta struct {
			Code string `json:"code"`
		} `json:"meta"`
	}{}
	if err := jsoniter.Unmarshal(resp, &r); err != nil {
		return
	}

	// 记录返回日志
	switch r.Meta.Code {
	case "SUCCESS":
		logx.Zap().Infow("api_response",
			"request_id", requestId,
			"url", requestUrl,
			"body", string(resp),
			"user_agent", userAgent,
			"client_ip", clientIp,
			"response_in_ms", timeDiff,
		)
	case "ERROR":
		logx.Zap().Errorw("api_response",
			"request_id", requestId,
			"url", requestUrl,
			"body", string(resp),
			"user_agent", userAgent,
			"client_ip", clientIp,
			"response_in_ms", timeDiff,
		)
	case "WARN":
		logx.Zap().Warnw("api_response",
			"request_id", requestId,
			"url", requestUrl,
			"body", string(resp),
			"user_agent", userAgent,
			"client_ip", clientIp,
			"response_in_ms", timeDiff,
		)
	default:
		logx.Zap().Infow("api_response",
			"request_id", requestId,
			"url", requestUrl,
			"body", string(resp),
			"user_agent", userAgent,
			"client_ip", clientIp,
			"response_in_ms", timeDiff,
		)
	}
}
